/*
 * @Author: 姚潘涛
 * @Date: 2024-06-04 09:21:05
 * @LastEditors: 姚潘涛
 * @LastEditTime: 2024-10-10 19:14:53
 * @Description:
 *
 * Copyright (c) 2024 by pandaman, All Rights Reserved.
 */
#include "xict_common/log/log.h"
#include "xict_common/perception/common/cuda_utils.h"
#include "xict_common/perception/common/logging.h"
#include "xict_common/perception/common/model.h"
#include "xict_common/perception/common/utils.h"
#include "xict_common/perception/detection/postprocess.h"
#include "xict_common/perception/detection/preprocess.h"
#include "xict_common/perception/detection/yolo_detect.h"
#include <chrono>
#include <cmath>
#include <iostream>

#define USE_IMG_DIR false
#define USE_CAMERA true
#define SERIALIZE_ENGINE false

using namespace nvinfer1;

static Logger gLogger;
const static int kOutputSize =
    kMaxNumOutputBbox * sizeof(xict_common::utils::Detection) / sizeof(float) +
    1;

namespace xict_common {
    namespace perception {
        namespace detection {
            class Yolov5TrtDet {
            public:
                Yolov5TrtDet(std::string engine);
                ~Yolov5TrtDet();
                void PrepareBuffers(ICudaEngine* engine,
                                    float** gpu_input_buffer,
                                    float** gpu_output_buffer,
                                    float** cpu_output_buffer);
                void Infer(IExecutionContext& context, cudaStream_t& stream,
                           void** gpu_buffers, float* output, int batchsize);
                void SerializeEngine(unsigned int max_batchsize, bool is_p6,
                                     float gd, float gw, std::string wts_name,
                                     std::string engine_name);
                void DeserializeEngine(std::string& engine_name,
                                       IRuntime** runtime, ICudaEngine** engine,
                                       IExecutionContext** context);
                void Detect(std::vector<cv::Mat> imgs,
                            std::vector<cv::Mat>& out_imgs);
                void DetectImgDir(std::string dir);
                void DetectCamera(cv::Mat img, cv::Mat& result_img);
                struct_yolo_result DetectOutputBbox(cv::Mat img);
                void DetectBboxBatch(std::vector<cv::Mat> imgs,
                                     std::vector<struct_yolo_result>& results);

            private:
                std::string wts_name_    = "";
                std::string engine_name_ = "";
                std::string net_         = "";
                bool is_p6_              = false;
                float gd_ = 0.0f, gw_ = 0.0f;
                std::string img_dir_;

                cudaStream_t stream_;

                float* gpu_buffers_[2];
                float* cpu_output_buffer_   = nullptr;
                IRuntime* runtime_          = nullptr;
                ICudaEngine* engine_        = nullptr;
                IExecutionContext* context_ = nullptr;
            };
        }   // namespace detection
    }       // namespace perception
}   // namespace xict_common